import math
import os

from PIL import Image
import numpy as np
from glob import glob
import os.path as osp
from util.vgrid_util import IQMsg
from tqdm import tqdm
import random


def makesure_dir(path):
    if not osp.exists(path):
        os.makedirs(path)


def read_nth_image(file_path, n, _width, _height, dtype=np.uint16):
    bytes_per_pixel = np.dtype(dtype).itemsize
    image_size_bytes = _width * _height * bytes_per_pixel
    # 计算第 n 个图像的起始位置
    offset = n * image_size_bytes
    # 读取目标图像
    image = np.fromfile(file_path, dtype=dtype, count=_width * _height, offset=offset)
    image = image.reshape((_height, _width))  # 调整形状为 (height, width)
    return image


if __name__ == '__main__':
    input_base_path = r"G:\不同部位数据集对\胸部-世纪坛\VGRID_B54_NOSNR"
    label_base_path = r"G:\不同部位数据集对\胸部-世纪坛\VGRID+SUB+SNR+HDR"
    save_path = r"G:\不同部位数据集对\HDRDataSet\胸部-世纪坛"
    sample_rate = 8

    input_save_path = osp.join(save_path, "input")
    label_save_path = osp.join(save_path, "label")
    _ = [makesure_dir(input_save_path), makesure_dir(label_save_path)]

    input_img_name = [osp.basename(name) for name in glob(osp.join(input_base_path, "*.raw"))]
    label_img_name = [osp.basename(name) for name in glob(osp.join(label_base_path, "*.raw"))]
    all_name_list = list(set(input_img_name + label_img_name))

    tqdm_bar = tqdm(all_name_list, total=len(all_name_list))
    iq_msg = IQMsg()
    for img_name in tqdm_bar:
        base_name = img_name.split(".")[0]
        input_path = osp.join(input_base_path, img_name)
        label_path = osp.join(label_base_path, img_name)
        iq_path = osp.join(input_base_path, img_name.split(".")[0] + ".iq")

        iq_msg.parse(iq_path)
        height, width, frame_count = iq_msg.height, iq_msg.width, iq_msg.Count

        used_frame_count = max(1, math.ceil(frame_count/sample_rate))
        used_frame_id = random.sample(list(range(frame_count)), used_frame_count)

        for frame_id in used_frame_id:
            input_img = read_nth_image(input_path, frame_id, width, height, dtype=np.uint16)
            label_img = read_nth_image(label_path, frame_id, width, height, dtype=np.uint16)
            Image.fromarray(input_img).save(osp.join(input_save_path, base_name + f"_frame{frame_id}.png"))
            Image.fromarray(label_img).save(osp.join(label_save_path, base_name + f"_frame{frame_id}.png"))







